package com.ahpu.lyf.huffmancode;

import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.HashMap;

public class HuffmanCode {
    public static void main(String[] args) {
        String s = "i like like like java do you like a java";
        byte[] bytes = s.getBytes(StandardCharsets.UTF_8);
        System.out.println(Arrays.toString(bytes));

        Huffman huffman = new Huffman();
        byte[] huffmanZipBytes = huffman.huffmanZip(bytes);
        System.out.println("压缩前长度:" + s.length());
        System.out.println("压缩后bytes数组：" + Arrays.toString(huffmanZipBytes) + "\t压缩后长度：" + huffmanZipBytes.length);

        System.out.println("~~~~~~~~~~~~解码后的字符串");
        byte[] decodeBytes = huffman.decode(huffman.huffmanCodes, huffmanZipBytes);
        System.out.println(new String(decodeBytes));


        /*      分步调用
        Node root = huffman.createHuffmanTree(bytes);
        huffman.orderTree(root);

        System.out.println("~~~~~~~~~~~~~HunffmanCodes");
        HashMap<Byte, String> huffmanCodes = huffman.getHuffmanCodes(root);
        System.out.println(huffmanCodes);

        huffmanCodes.forEach((k,v)->System.out.print((char) k.intValue()+" = "+v+"\t"));
        System.out.println("\n~~~~~~~~~~~~~~~赫夫曼编码目标字符串s");
        byte[] zip = huffman.zip(bytes, huffmanCodes);
        System.out.println(Arrays.toString(zip));*/
    }
}


class Huffman {
    HashMap<Byte, String> huffmanCodes = new HashMap<>();
    StringBuilder tempCode = new StringBuilder();

    /**
     * huffman压缩
     *
     * @param bytes 原始的bytes数组
     * @return 压缩后的bytes数组
     */
    public byte[] huffmanZip(byte[] bytes) {
        //得到赫夫曼树
        Node huffmanTree = createHuffmanTree(bytes);
        //得到赫夫曼编码
        HashMap<Byte, String> huffmanCodes = getHuffmanCodes(huffmanTree);
        //返回赫夫曼压缩后的bytes数组
        return zip(bytes, huffmanCodes);
    }


    /**
     * 构建赫夫曼树
     *
     * @param bytes 字符串的baty数组
     * @return 返回赫夫曼树
     */
    public Node createHuffmanTree(byte[] bytes) {
        //创建一个map保存数据key：字符，value：权重
        HashMap<Byte, Integer> weightMap = new HashMap<>();
        for (byte b : bytes) {
            //字符存在则权重加一，不存在则保存并赋权重为1
            weightMap.merge(b, 1, Integer::sum);
        }

        //创建一个list保存节点
        ArrayList<Node> nodes = new ArrayList<>();

        //将每个字符转化为一个最简树（节点，左右为null）
        weightMap.forEach((k, v) -> nodes.add(new Node(k, v)));

        //构建赫夫曼树
        while (nodes.size() > 1) {
            //给nodes升序排序
            Collections.sort(nodes);
            //取权重最小的两个节点，构成一个新树
            Node leftNode = nodes.get(0);
            Node rightNode = nodes.get(1);
            Node parentNode = new Node(null, leftNode.weight + rightNode.weight);
            parentNode.left = leftNode;
            parentNode.right = rightNode;
            //移除这两个树，并将新树添加到list
            nodes.remove(leftNode);
            nodes.remove(rightNode);
            nodes.add(parentNode);
        }

        //最后留在list的树就是最终的哈夫曼树
        return nodes.get(0);
    }

    /**
     * 遍历赫夫曼树
     *
     * @param root 赫夫曼树的根节点
     */
    public void orderTree(Node root) {
        if (root != null) {
            System.out.println("~~~~~~~~~~~~~前序遍历");
            root.preOrder(root);
            System.out.println("~~~~~~~~~~~~~中序遍历");
            root.infixOrder(root);
        }
    }

    /**
     * 对外的获取赫夫曼编码的方法
     *
     * @param root 根节点
     * @return 返回赫夫曼编码的map
     */
    public HashMap<Byte, String> getHuffmanCodes(Node root) {
        getHuffmanCodes(root, "", tempCode);
        return huffmanCodes;
    }

    /**
     * 得到哈夫曼编码的方法
     *
     * @param node     当前节点
     * @param code     规定走左子树路径为0，走右边路径子树为1
     * @param tempCode 外部变量用来保存递归时的路径
     */
    private void getHuffmanCodes(Node node, String code, StringBuilder tempCode) {
        //加载上次递归保存的路径
        StringBuilder resCode = new StringBuilder(tempCode);
        //将当前走的路径添加进去
        resCode.append(code);
        if (node != null) {
            //如果是非叶子节点
            if (node.data == null) {
                //向左递归
                getHuffmanCodes(node.left, "0", resCode);
                //向右递归
                getHuffmanCodes(node.right, "1", resCode);
                //是叶子节点，添加编码和路径到map
            } else {
                huffmanCodes.put(node.data, resCode.toString());
            }
        }
    }

    /**
     * huffman压缩
     *
     * @param bytes        要压缩的bytes数组
     * @param huffmanCodes 赫夫曼编码
     * @return 压缩后的bytes数组
     */
    public byte[] zip(byte[] bytes, HashMap<Byte, String> huffmanCodes) {
//        拼接成一个字符串
        StringBuilder stringBuilder = new StringBuilder();
        for (byte aByte : bytes) {
            stringBuilder.append(huffmanCodes.get(aByte));
        }
        //1010100010111111110010001011111111001000101111111100100101001101110001110000011011101000111100101000101111111100110001001010011011100
//        System.out.println(stringBuilder);
        //133
//        System.out.println(stringBuilder.length());
        //为什么要加1后取模8？？？因为如果他正好是8的倍数多1时，刚好结果就是多少组byte（8位）
        int len = (stringBuilder.length() + 7) / 8;

        byte[] huffmanCodeBytes = new byte[len];
        int index = 0;
        //8位一取得到一个字节
        for (int i = 0; i < stringBuilder.length(); i += 8) {
            String tempStr;
            //最后少于8位
            if (i + 8 > stringBuilder.length()) {
                //取后面所有的
                tempStr = stringBuilder.substring(i);
            } else {
                //取8位
                tempStr = stringBuilder.substring(i, i + 8);
            }
            //将8位的字符串按照2进制转换为int再强转为byte
//            System.out.println(tempStr);
            //存到byte是补码，源码-1再取反，最高位是符号位1表示负数
            huffmanCodeBytes[index] = (byte) Integer.parseInt(tempStr, 2);
            index++;
        }
        return huffmanCodeBytes;
    }


    /**
     * 将一个byte转换成8位字符串
     *
     * @param flag 是否需要补高位
     * @param b    要转换的字节
     * @return 转换后的8位字符串
     */
    public String byteToBitString(boolean flag, byte b) {
        //使用int类型变量保存b，便于转换
        int temp = b;
        //是正数需要补高位
        if (flag) {
            //temp与256按位取与，补高位用,256是1 0000 0000，与这个取与补位其他为不变
            temp |= 256;
        }
        String binaryString = Integer.toBinaryString(temp);
//        System.out.println(binaryString);
        if (flag)
            //int补高位后是32位，只要最后8位即可
            return binaryString.substring(binaryString.length() - 8);
        return binaryString;
    }

    /**
     * 解码
     *
     * @param huffmanCodes 哈夫曼编码
     * @param huffmanBytes 编码后的字节数组
     * @return 原字符串对应的byte数组
     */
    public byte[] decode(HashMap<Byte, String> huffmanCodes, byte[] huffmanBytes) {
        StringBuilder huffmanSb = new StringBuilder();
        for (int i = 0; i < huffmanBytes.length; i++) {
            //判断是不是最后一位，是最后一位不用补高位
            boolean flag = (i == huffmanBytes.length - 1);
            huffmanSb.append(byteToBitString(!flag, huffmanBytes[i]));
        }
//        System.out.println(huffmanSb);
        //赫夫曼编码表反转
        HashMap<String,Byte> decodeMap=new HashMap<>();
        huffmanCodes.forEach((k,v)->decodeMap.put(v,k));

        //保存byte
        ArrayList<Byte> byteList = new ArrayList<>();
        //暴力匹配
        for (int i = 0; i < huffmanSb.length();) {
            int count=1;
            boolean flag=true;
            Byte b=null;
            while (flag){
                //从1位开始匹配
                String key= huffmanSb.substring(i,i+count);
                //匹配解码表的结果
                b = decodeMap.get(key);
                //如果为null，说明还没有匹配到
                if (b==null)
                    //匹配的长度加一
                    count++;
                else
                    //匹配到了跳出循环
                    flag=false;
            }
            //将结果放到列表
            byteList.add(b);
            //下一位后移匹配过的位数
            i+=count;
        }
        //保存字节的数组
        byte[] bytes=new byte[byteList.size()];
        //将字节列表的值放到字节数组
        for (int i = 0; i < bytes.length; i++) {
            bytes[i]= byteList.get(i);
        }
        return bytes;
    }

}


class Node implements Comparable<Node> {
    //用来存放数据的
    Byte data;
    //字符出现的次数作为权重
    Integer weight;
    Node left;
    Node right;

    /**
     * 前序遍历
     *
     * @param node 当前节点
     */
    public void preOrder(Node node) {
        System.out.println(node);
        if (node.left != null)
            preOrder(node.left);
        if (node.right != null)
            preOrder(node.right);
    }

    /**
     * 中序遍历
     *
     * @param node 当前节点
     */
    public void infixOrder(Node node) {
        if (node.left != null)
            infixOrder(node.left);
        System.out.println(node);
        if (node.right != null)
            infixOrder(node.right);
    }

    public Node(Byte data, Integer weight) {
        this.data = data;
        this.weight = weight;
    }

    @Override
    public String toString() {
        return "Node{" +
                "data=" + data +
                ", weight=" + weight +
                '}';
    }

    @Override
    public int compareTo(Node o) {
        //升序排列
        return this.weight - o.weight;
    }
}
